import os
from scipy.io import mmread
import numpy as np

def judge_diagonal_matrix(file_path):
    """统计 .mtx 文件中对角元素的数量"""
    matrix = mmread(file_path).tocoo()
    diag_count = np.sum(matrix.row == matrix.col)
    return 1 if diag_count == matrix.nnz else 0, matrix.nnz

def count_diagonal_entries_in_directory(directory):
    num_diag_mtx = 0
    num_mtx = 0
    nnzs = []
    for root, dirs, files in os.walk(directory):
        for file in files:
            if file.endswith(".mtx"):
                num_mtx += 1
                path = os.path.join(root, file)
                try:
                    isdiag, nnz = judge_diagonal_matrix(path)
                    if isdiag:
                        num_diag_mtx+=1
                    nnzs.append((nnz, file))
                except Exception as e:
                    print(f"Error processing {path}: {e}")
    
    return num_diag_mtx, num_mtx, nnzs

# 设置你的目录路径
directory_path = "/home/yida/demo/analyze/PanguLU_next_exportdiag/examples/solverchallenge25_03_A.mtx.diagblks_500"

print(directory_path)
diag_cnt, mtx_cnt, nnzs = count_diagonal_entries_in_directory(directory_path)

print(f"对角矩阵统计：{diag_cnt} / {mtx_cnt}")
# nnzs.sort()
# length = len(nnzs)
# nstep = 10
# steplen = length // nstep
# for i in range(nstep):
#     print(f"{nnzs[i * steplen]/2500}% --- ", end="")
# print(f"{nnzs[-1]/2500}%")

sparsity = [0 for i in range(11)]

for i in range(len(nnzs)):
    sparsity[int(nnzs[i][0] // (500 * 500 / 10))] += 1

print(sparsity)

nnz_csv = open("03_A.csv", "w")
nnz_csv.write("level, nnz\n")
for i in range(len(nnzs)):
    nnz_csv.write(f"{int(nnzs[i][1].split(".")[0])}, {nnzs[i][0]}\n")